import base64
import logging
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from .models import qkd_keys

logger = logging.getLogger(__name__)


class AESCipher:
    #AES-128加密/解密

    @staticmethod
    def encrypt(raw_data, key_or_username):
        try:
            # 判断是用户名还是直接的密钥
            if len(key_or_username) == 32 and all(c in '0123456789abcdef' for c in key_or_username.lower()):
                # 看起来是直接的密钥
                hex_key = key_or_username
            elif key_or_username in qkd_keys:
                # 这是用户名，获取密钥
                hex_key = qkd_keys[key_or_username]
            else:
                logger.error(f"无效的密钥或用户名: {key_or_username}")
                return None

            # 使用密钥加密
            key = bytes.fromhex(hex_key)

            # 创建AES-128-CBC加密器
            cipher = AES.new(key, AES.MODE_CBC)

            # 加密数据并进行填充
            padded_data = pad(raw_data.encode('utf-8'), AES.block_size)
            encrypted_data = cipher.encrypt(padded_data)

            # 将IV和加密数据合并并进行base64编码
            result = base64.b64encode(cipher.iv + encrypted_data).decode('utf-8')

            return result
        except Exception as e:
            logger.error(f"AES加密失败: {str(e)}")
            return None

    @staticmethod
    def decrypt(encrypted_data, key_or_username):
        try:
            # 判断是用户名还是直接的密钥
            if len(key_or_username) == 32 and all(c in '0123456789abcdef' for c in key_or_username.lower()):
                # 看起来是直接的密钥
                hex_key = key_or_username
            elif key_or_username in qkd_keys:
                # 这是用户名，获取密钥
                hex_key = qkd_keys[key_or_username]
            else:
                logger.error(f"无效的密钥或用户名: {key_or_username}")
                return None

            # 使用密钥解密
            key = bytes.fromhex(hex_key)

            # 解码base64数据
            raw = base64.b64decode(encrypted_data)

            # 分离IV和加密数据
            iv = raw[:AES.block_size]
            encrypted_data = raw[AES.block_size:]

            # 创建解密器
            cipher = AES.new(key, AES.MODE_CBC, iv)

            # 解密数据并去除填充
            decrypted_data = unpad(cipher.decrypt(encrypted_data), AES.block_size).decode('utf-8')

            return decrypted_data
        except Exception as e:
            logger.error(f"AES解密失败: {str(e)}")
            return None